import torch
import torch.nn as nn
import torch.nn.functional as F

from nets.resnet import MyResNet

def get_img_output_length(width, height):
    def get_output_length(input_length):
        # input_length += 6
        filter_sizes = [2, 2, 2, 2, 2]
        padding = [0, 0, 0, 0, 0]
        stride = 2
        for i in range(5):
            input_length = (input_length + 2 * padding[i] - filter_sizes[i]) // stride + 1
        return input_length
    return get_output_length(width) * get_output_length(height) 
    
class Siamese(nn.Module):
    def __init__(self, input_shape, pretrained=True):
        super(Siamese, self).__init__()
        self.ResNet = MyResNet(pretrained)
        # del self.ResNet.avgpool
        # del self.vgg.classifier
        
        # flat_shape = 512 * get_img_output_length(input_shape[1], input_shape[0])
        # print(flat_shape)
        self.fully_connect1 = torch.nn.Linear(2048*7*7, 512)
        self.fully_connect2 = torch.nn.Linear(512, 2)

    def forward(self, x):


        #------------------------------------------#
        #   我们将两个输入传入到主干特征提取网络
        # 这里x1为left x2为right
        #------------------------------------------#

        # x1 = F.interpolate(x1, size=(224,224), mode='nearest')
        # x2 = F.interpolate(x2, size=(224, 224), mode='nearest')
        # n1 = x1[0]
        # n = self.vgg.features(n1)

        # from torchsummary import summary
        # summary(self.ResNet.features, input_size=[(3,224, 224)],batch_size=8)

        x = self.ResNet(x)#
        # x2 = self.ResNet(x2)
        #-------------------------#
        #   相减取绝对值
        #-------------------------#
        x = torch.flatten(x, 1)
        # x2 = torch.flatten(x2, 1)
        # x = torch.abs(x1 - x2)
        #-------------------------#
        #   进行两次全连接
        #-------------------------#

        x = self.fully_connect1(x)
        x = self.fully_connect2(x)
        return x

if __name__ == '__main__':
    # model = torchvision.models.resnet50()
    model = Siamese()
    # print(model)
    from torchsummary import summary
    summary(model, input_size=[(3, 224, 224)], batch_size=2)

    # input = torch.randn(1, 3, 224, 224)
    # out = model(input)
    # print(out.shape)

